import pandas as pd
import json
import re
from pathlib import Path


def extract_water_system_numbers_from_text(text):
    """从水系描述文本中提取数字信息"""
    if pd.isna(text) or not isinstance(text, str):
        return None

    # 提取长度的模式
    length_patterns = [
        r'长\s*(\d+(?:\.\d+)?)\s*余?里',
        r'长\s*(\d+(?:\.\d+)?)\s*余?公里',
        r'全长\s*(\d+(?:\.\d+)?)\s*余?里',
        r'全长\s*(\d+(?:\.\d+)?)\s*余?公里',
        r'(\d+(?:\.\d+)?)\s*余?里',
        r'(\d+(?:\.\d+)?)\s*余?公里',
        r'东西长约\s*(\d+(?:\.\d+)?)\s*公里',
        r'南北宽约\s*(\d+(?:\.\d+)?)\s*公里',
    ]

    # 提取宽度的模式
    width_patterns = [
        r'宽\s*(\d+(?:\.\d+)?)\s*余?里',
        r'宽\s*(\d+(?:\.\d+)?)\s*余?公里',
        r'宽度\s*(\d+(?:\.\d+)?)\s*余?里',
        r'宽度\s*(\d+(?:\.\d+)?)\s*余?公里',
        r'东西宽\s*(\d+(?:\.\d+)?)\s*余?里',
        r'南北宽\s*(\d+(?:\.\d+)?)\s*余?里',
    ]

    # 提取年份的模式
    year_patterns = [
        r'(\d{3,4})\s*年',
        r'公元\s*(\d{3,4})',
        r'开皇\s*(\d+)',
        r'大业\s*(\d+)',
        r'至元\s*(\d+)',
        r'大定\s*(\d+)',
        r'泰和\s*(\d+)',
        r'光绪\s*(\d+)',
        r'宣统\s*(\d+)',
        r'民国\s*(\d+)',
    ]

    length = None
    width = None
    year = None

    # 提取长度
    for pattern in length_patterns:
        match = re.search(pattern, text)
        if match:
            try:
                num = float(match.group(1))
                length = num
                break
            except:
                pass

    # 提取宽度
    for pattern in width_patterns:
        match = re.search(pattern, text)
        if match:
            try:
                num = float(match.group(1))
                width = num
                break
            except:
                pass

    # 提取年份
    for pattern in year_patterns:
        match = re.search(pattern, text)
        if match:
            try:
                year_val = int(match.group(1))
                # 处理年号转公元年的情况
                if year_val < 100:  # 可能是年号年份
                    # 根据朝代大致转换
                    if '开皇' in text and year_val <= 20:
                        year = 580 + year_val  # 开皇元年为581年
                    elif '大业' in text and year_val <= 14:
                        year = 604 + year_val  # 大业元年为605年
                    elif '至元' in text:
                        year = 1264 + year_val  # 至元元年为1264年
                    elif '大定' in text:
                        year = 1161 + year_val  # 大定元年为1161年
                    elif '泰和' in text:
                        year = 1201 + year_val  # 泰和元年为1201年
                    elif '光绪' in text:
                        year = 1875 + year_val  # 光绪元年为1875年
                    elif '宣统' in text:
                        year = 1909 + year_val  # 宣统元年为1909年
                    elif '民国' in text:
                        year = 1911 + year_val  # 民国元年为1912年
                    else:
                        continue
                else:
                    year = year_val
                break
            except:
                pass

    return {'length': length, 'width': width, 'year': year}


def extract_water_system_data_from_excel(file_path, sheet_name='总水系'):
    """从Excel文件中提取水系数据"""
    try:
        # 读取Excel文件
        df = pd.read_excel(file_path, sheet_name=sheet_name)

        print(f"数据框形状: {df.shape}")
        print(f"列名: {df.columns.tolist()}")

        results = []

        for idx, row in df.iterrows():
            # 获取时期（朝代）
            period = row.iloc[0] if not pd.isna(row.iloc[0]) else None

            # 获取河流/水系名称
            river_name = row.iloc[3] if len(row) > 3 and not pd.isna(row.iloc[3]) else None

            # 获取水系总体描述的原文
            content = row.iloc[4] if len(row) > 4 and not pd.isna(row.iloc[4]) else None

            if period and content:
                # 从文本中提取数字信息
                extracted = extract_water_system_numbers_from_text(str(content))

                # 创建数据记录
                record = {
                    'category': str(period),
                    'river_name': str(river_name) if river_name else None,
                    'year': extracted['year'],
                    'length': extracted['length'],
                    'width': extracted['width']
                }

                # 只添加至少有一个数值字段的记录
                if any([record['year'], record['length'], record['width']]):
                    results.append(record)

        return results

    except Exception as e:
        print(f"读取文件时出错: {e}")
        import traceback
        traceback.print_exc()
        return []


def main():
    # 输入文件路径
    input_file = './data/02水系 - 总数据和各朝代数据.xlsx'

    # 输出文件路径
    output_file = './outputs/water_system_data.json'

    print("开始提取水系数据...")

    # 提取水系数据
    data = extract_water_system_data_from_excel(input_file, sheet_name='总水系')

    print(f"从'总水系' sheet提取了 {len(data)} 条记录")

    # 保存为JSON文件
    Path(output_file).parent.mkdir(parents=True, exist_ok=True)
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)

    print(f"数据已保存到: {output_file}")

    # 显示前几条数据作为示例
    print("\n前10条数据示例:")
    for i, record in enumerate(data[:10], 1):
        print(f"\n记录 {i}:")
        print(f"  时期(category): {record['category']}")
        print(f"  河流名称(river_name): {record['river_name']}")
        print(f"  年份(year): {record['year']}")
        print(f"  长度(length): {record['length']} 里" if record['length'] else f"  长度(length): {record['length']}")
        print(f"  宽度(width): {record['width']} 里" if record['width'] else f"  宽度(width): {record['width']}")

    print(f"\n总共提取了 {len(data)} 条有效记录")

    # 统计信息
    years = [r['year'] for r in data if r['year']]
    lengths = [r['length'] for r in data if r['length']]
    widths = [r['width'] for r in data if r['width']]

    print(f"\n统计信息:")
    print(f"  包含年份的记录: {len(years)} 条")
    print(f"  包含长度的记录: {len(lengths)} 条")
    print(f"  包含宽度的记录: {len(widths)} 条")

    if lengths:
        print(f"  最长河流: {max(lengths)} 里")
        print(f"  平均长度: {sum(lengths) / len(lengths):.2f} 里")

    if widths:
        print(f"  最宽河流: {max(widths)} 里")
        print(f"  平均宽度: {sum(widths) / len(widths):.2f} 里")

    return data


if __name__ == '__main__':
    data = main()
